//! WebSocket step implementation

use crate::error::{Error, Result};
use crate::models::{
    config::StepConfig,
    step::{IStep, StepType, StepResult, StepContext},
    Value, ValueMap,
};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::time::{Duration, Instant};
use tokio_tungstenite::{connect_async, tungstenite::Message};
use uuid::Uuid;

/// WebSocket action types
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum WebSocketAction {
    Connect,
    Send,
    Receive,
    Close,
}

/// WebSocket step configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WebSocketConfig {
    /// WebSocket action to perform
    pub action: WebSocketAction,

    /// WebSocket URL (for connect action)
    pub url: Option<String>,

    /// Message to send (for send action)
    pub message: Option<Value>,

    /// Expected message pattern (for receive action)
    pub expect: Option<Value>,

    /// Timeout for the operation
    pub timeout: Option<u64>,

    /// Connection ID for reusing connections
    pub connection_id: Option<String>,
}

/// WebSocket step implementation
#[derive(Debug, Clone)]
pub struct WebSocketStep {
    id: Uuid,
    config: StepConfig,
    websocket_config: WebSocketConfig,
}

impl WebSocketStep {
    /// Create a new WebSocket step
    pub fn new(config: StepConfig, websocket_config: WebSocketConfig) -> Self {
        Self {
            id: Uuid::new_v4(),
            config,
            websocket_config,
        }
    }

    /// Create a connect step
    pub fn connect(url: String, connection_id: Option<String>) -> Self {
        let config = StepConfig {
            step_name: "WebSocket Connect".to_string(),
            ..Default::default()
        };

        let websocket_config = WebSocketConfig {
            action: WebSocketAction::Connect,
            url: Some(url),
            message: None,
            expect: None,
            timeout: Some(30), // 30 seconds default timeout
            connection_id,
        };

        Self::new(config, websocket_config)
    }

    /// Create a send step
    pub fn send(message: Value, connection_id: Option<String>) -> Self {
        let config = StepConfig {
            step_name: "WebSocket Send".to_string(),
            ..Default::default()
        };

        let websocket_config = WebSocketConfig {
            action: WebSocketAction::Send,
            url: None,
            message: Some(message),
            expect: None,
            timeout: Some(10), // 10 seconds default timeout
            connection_id,
        };

        Self::new(config, websocket_config)
    }

    /// Create a receive step
    pub fn receive(expect: Option<Value>, connection_id: Option<String>) -> Self {
        let config = StepConfig {
            step_name: "WebSocket Receive".to_string(),
            ..Default::default()
        };

        let websocket_config = WebSocketConfig {
            action: WebSocketAction::Receive,
            url: None,
            message: None,
            expect,
            timeout: Some(10), // 10 seconds default timeout
            connection_id,
        };

        Self::new(config, websocket_config)
    }

    /// Create a close step
    pub fn close(connection_id: Option<String>) -> Self {
        let config = StepConfig {
            step_name: "WebSocket Close".to_string(),
            ..Default::default()
        };

        let websocket_config = WebSocketConfig {
            action: WebSocketAction::Close,
            url: None,
            message: None,
            expect: None,
            timeout: Some(5), // 5 seconds default timeout
            connection_id,
        };

        Self::new(config, websocket_config)
    }
}

#[async_trait]
impl IStep for WebSocketStep {
    fn name(&self) -> &str {
        &self.config.step_name
    }

    fn step_type(&self) -> StepType {
        StepType::WebSocket
    }

    fn config(&self) -> &StepConfig {
        &self.config
    }

    fn config_mut(&mut self) -> &mut StepConfig {
        &mut self.config
    }

    async fn execute(&self, context: &mut StepContext) -> Result<StepResult> {
        let start_time = Instant::now();
        let mut step_result = StepResult::new(self.id, self.name().to_string(), self.step_type());

        // Execute based on action type
        match self.websocket_config.action {
            WebSocketAction::Connect => {
                self.execute_connect(context, &mut step_result).await?;
            }
            WebSocketAction::Send => {
                self.execute_send(context, &mut step_result).await?;
            }
            WebSocketAction::Receive => {
                self.execute_receive(context, &mut step_result).await?;
            }
            WebSocketAction::Close => {
                self.execute_close(context, &mut step_result).await?;
            }
        }

        step_result.elapsed = start_time.elapsed().as_millis() as i64;
        Ok(step_result)
    }

    fn id(&self) -> Uuid {
        self.id
    }

    fn clone_step(&self) -> Box<dyn IStep> {
        Box::new(self.clone())
    }
}

impl WebSocketStep {
    /// Execute connect action
    async fn execute_connect(&self, context: &mut StepContext, step_result: &mut StepResult) -> Result<()> {
        let url = self.websocket_config.url.as_ref()
            .ok_or_else(|| Error::Internal("WebSocket URL is required for connect action".to_string()))?;

        let timeout = Duration::from_secs(self.websocket_config.timeout.unwrap_or(30));

        // For now, we'll simulate the connection since we don't have a full WebSocket manager
        // In a real implementation, this would establish and store the connection

        let mut response_data = ValueMap::new();
        response_data.insert("action".to_string(), Value::String("connect".to_string()));
        response_data.insert("url".to_string(), Value::String(url.clone()));
        response_data.insert("status".to_string(), Value::String("connected".to_string()));

        if let Some(conn_id) = &self.websocket_config.connection_id {
            response_data.insert("connection_id".to_string(), Value::String(conn_id.clone()));
            // Store connection in context for later use
            context.variables.insert(
                format!("websocket_connection_{}", conn_id),
                Value::String(url.clone())
            );
        }

        step_result.data = Some(Value::Object(response_data.into_iter().collect()));
        step_result.success = true;

        Ok(())
    }

    /// Execute send action
    async fn execute_send(&self, context: &mut StepContext, step_result: &mut StepResult) -> Result<()> {
        let message = self.websocket_config.message.as_ref()
            .ok_or_else(|| Error::Internal("Message is required for send action".to_string()))?;

        // Verify connection exists if connection_id is specified
        if let Some(conn_id) = &self.websocket_config.connection_id {
            let conn_key = format!("websocket_connection_{}", conn_id);
            if !context.variables.contains_key(&conn_key) {
                return Err(Error::Internal(format!("WebSocket connection '{}' not found", conn_id)));
            }
        }

        let mut response_data = ValueMap::new();
        response_data.insert("action".to_string(), Value::String("send".to_string()));
        response_data.insert("message".to_string(), message.clone());
        response_data.insert("status".to_string(), Value::String("sent".to_string()));

        if let Some(conn_id) = &self.websocket_config.connection_id {
            response_data.insert("connection_id".to_string(), Value::String(conn_id.clone()));
        }

        step_result.data = Some(Value::Object(response_data.into_iter().collect()));
        step_result.success = true;

        Ok(())
    }

    /// Execute receive action
    async fn execute_receive(&self, context: &mut StepContext, step_result: &mut StepResult) -> Result<()> {
        // Verify connection exists if connection_id is specified
        if let Some(conn_id) = &self.websocket_config.connection_id {
            let conn_key = format!("websocket_connection_{}", conn_id);
            if !context.variables.contains_key(&conn_key) {
                return Err(Error::Internal(format!("WebSocket connection '{}' not found", conn_id)));
            }
        }

        // For now, simulate receiving a message
        let received_message = Value::String("Simulated received message".to_string());

        let mut response_data = ValueMap::new();
        response_data.insert("action".to_string(), Value::String("receive".to_string()));
        response_data.insert("message".to_string(), received_message.clone());
        response_data.insert("status".to_string(), Value::String("received".to_string()));

        if let Some(conn_id) = &self.websocket_config.connection_id {
            response_data.insert("connection_id".to_string(), Value::String(conn_id.clone()));
        }

        // Check if message matches expected pattern
        if let Some(expected) = &self.websocket_config.expect {
            let matches = match (expected, &received_message) {
                (Value::String(exp), Value::String(recv)) => exp == recv,
                _ => expected == &received_message,
            };

            response_data.insert("matches_expected".to_string(), Value::Bool(matches));

            if !matches {
                step_result.success = false;
                step_result.error = Some("Received message does not match expected pattern".to_string());
            } else {
                step_result.success = true;
            }
        } else {
            step_result.success = true;
        }

        step_result.data = Some(Value::Object(response_data.into_iter().collect()));

        Ok(())
    }

    /// Execute close action
    async fn execute_close(&self, context: &mut StepContext, step_result: &mut StepResult) -> Result<()> {
        // Remove connection from context if connection_id is specified
        if let Some(conn_id) = &self.websocket_config.connection_id {
            let conn_key = format!("websocket_connection_{}", conn_id);
            context.variables.remove(&conn_key);
        }

        let mut response_data = ValueMap::new();
        response_data.insert("action".to_string(), Value::String("close".to_string()));
        response_data.insert("status".to_string(), Value::String("closed".to_string()));

        if let Some(conn_id) = &self.websocket_config.connection_id {
            response_data.insert("connection_id".to_string(), Value::String(conn_id.clone()));
        }

        step_result.data = Some(Value::Object(response_data.into_iter().collect()));
        step_result.success = true;

        Ok(())
    }
}
